
"""
智能问数查询改写
"""
from core.state import AnalyzeState
from langgraph.types import Command
from langgraph.config import get_stream_writer
from core.llm import rewrite_query_analyze_wyx
import concurrent.futures
from loguru import logger
from copy import deepcopy


async def wyx_rewrite_query_node(state: AnalyzeState):

    writer = get_stream_writer()
    writer("智能问数查询改写")
    logger.info("智能问数查询改写")

    schemas = state.get('schemas')
    query = state.get('query')
    rewrites = state.get('rewrites')
    rewrites.append(query)

    # 智能问数查询改写。
    userinfo = state.get('userinfo')

    short_schemas = []
    for schema in schemas:
        copy_schema = deepcopy(schema)
        copy_schema.pop('rows')
        copy_schema.pop('columns')
        short_schemas.append(copy_schema)

    rewrites = rewrite_query_analyze_wyx.invoke(userinfo.description, query, short_schemas)

    writer(rewrites)
    logger.info(f'改写的问题描述列表：{rewrites}')

    return Command(update={"rewrites": rewrites}, goto='biw_query')
